# encoding: utf-8

from typing import Callable
import pandas as pd
from vnpy.trader.object import Exchange, BarData, Interval
import math
from datetime import datetime
import os
 
class VnpyDataService:
    @classmethod
    def import_bar_data_from_file(cls, workspace_dir: str, file_path: str, process_callback: Callable):
        origin_cwd = os.getcwd()
        try:
            os.chdir(workspace_dir)
            df = pd.read_parquet(file_path)
            
            from vnpy.trader.utility import ZoneInfo
            from vnpy.trader.setting import SETTINGS
            from vnpy.trader.database import get_database
            database = get_database()

            # 获取数据库时区设置
            db_tz = ZoneInfo(SETTINGS["database.timezone"])

            total_rows = len(df)
            if total_rows == 0:
                process_callback(100)
                return

            processed_rows = 0
            batch_size = 10000
            
            grouped = df.groupby('symbol')

            for symbol, group_df in grouped:
                group_total_rows = len(group_df)
                num_batches = math.ceil(group_total_rows / batch_size)

                if num_batches == 0:
                    continue

                for i in range(num_batches):
                    start = i * batch_size
                    end = start + batch_size
                    batch_df = group_df.iloc[start:end]

                    bar_data_list = []
                    for row in batch_df.itertuples():
                        # 获取原始datetime
                        raw_datetime = getattr(row, "datetime")
                        
                        # 如果datetime没有时区信息，先本地化到数据库时区
                        if raw_datetime.tz is None:
                            localized_datetime = raw_datetime.tz_localize(db_tz)
                        else:
                            localized_datetime = raw_datetime
                        
                        # 转换为标准的Python datetime对象
                        python_datetime = localized_datetime.to_pydatetime()
                        
                        bar = BarData(
                            gateway_name='RQ',
                            symbol=getattr(row, "symbol"),
                            exchange=Exchange(getattr(row, "exchange")),
                            datetime=python_datetime,
                            open_price=getattr(row, "open_price"),
                            high_price=getattr(row, "high_price"),
                            low_price=getattr(row, "low_price"),
                            close_price=getattr(row, "close_price"),
                            interval=Interval(getattr(row, "interval")),
                            volume=getattr(row, "volume"),
                            turnover=getattr(row, "turnover"),
                            open_interest=getattr(row, "open_interest"),
                        )
                        bar_data_list.append(bar)
                    
                    database.save_bar_data(bar_data_list)
                    
                    processed_rows += len(batch_df)
                    progress = int(processed_rows / total_rows * 100)
                    process_callback(progress)
        finally:
            os.chdir(origin_cwd)

    @classmethod
    def import_contract_info_from_file(cls, file_path: str):
        """从文件导入合约信息"""
        df = pd.read_excel(file_path)
        
        required_columns = {'vnpy_underlying_symbol', 'contract_size', 'exchange', 'name'}
        if not required_columns.issubset(df.columns):
            missing_cols = required_columns - set(df.columns)
            raise ValueError(f"文件缺少必需的列: {', '.join(missing_cols)}")

        from dao.contract_info_dao import ContractInfoDao
        contract_info_dao = ContractInfoDao()
        contract_info_dao.save_contract_info(df)